# TODO: 导入必要的库和模块
import pickle
import gradio as gr
import numpy as np
from collections import Counter
from pinecone_example import load_pinecone



import matplotlib.pyplot as plt
from PIL import Image


# # TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
# with open('best_knn_model.pkl', 'rb') as model_file:
#     knn_model = pickle.load(model_file)

index = load_pinecone()


# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(stroke):
    # print(stroke)
    image = stroke['composite']
    # image.save('temp1.png')

    # 将图像数据转换为一维数组
    image_array = np.array(image)

    plt.imshow(image_array, cmap='gray')
    plt.axis('off')  # 关闭坐标轴
    plt.savefig('output_image.png')  # 保存图像到文件
    plt.close()

    # 将图像数据转换为一维数组并重新调整大小
    image_array = image_array.flatten().reshape(-1, 64)


    """
    # 使用模型进行预测
    prediction = knn_model.predict(image_array)
    # print(prediction)

    predicted_label = int(prediction.argmax())
    # print(predicted_label)
    return predicted_label
    """
    # 使用 Pinecone 进行查询
    results = index.query(vector=image_array[0].tolist(), top_k=11, include_metadata=True)

    # 提取 labels
    labels = [match['metadata']['label'] for match in results['matches']]

    # 使用投票机制决定最终的预测数字
    if labels:
        predicted_label = Counter(labels).most_common(1)[0][0]
    else:
        predicted_label = None  # 可以选择返回一个默认值或处理无匹配的情况

    return predicted_label

# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="绘制数字",crop_size=(8,8),type="pil"),
    outputs="label",
    title="数字识别",
    description="绘制数字并识别其数字"
)

# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch()